from __future__ import absolute_import, division, print_function, \
    unicode_literals
import sys
sys.path.append("..")
import argparse
import numpy as np
import os
import cv2
from nfs.aiflow.report import *
import uuid

# load data
def load_data(dir_path,if_gray):
    files = os.listdir(dir_path)
    files.sort(key=lambda x: int(x[:-4]))  # 倒着数第四位'.'为分界线，按照‘.’左边的数字从小到大排序

    x = []
    y = []
    for file in files:
        img_path = os.path.join(dir_path, file)
        # 默认是读入为彩色图，即使原图是灰度图也会复制成三个相同的通道变成彩色图
        # 第二个参数为0的时候读入为灰度图，即使原图是彩色图也会转成灰度图,灰度图少一个维度
        if if_gray:
            img = cv2.imread(img_path, 0)
        else:
            img = cv2.imread(img_path)
        label = file.split('.')[0].split('_')[-1]
        x.append(img)
        y.append(label)
    x = np.array(x)
    y = np.array(y)
    return x,y


# do data transform
def transform(output_dir, file_name,if_gray):
    x_train_name = uuid.uuid1().__str__()#'x_train.npy'
    x_test_name = uuid.uuid1().__str__()#'x_test.npy'
    y_train_name = uuid.uuid1().__str__()#'y_train.npy'
    y_test_name = uuid.uuid1().__str__()#'y_test.npy'
    if if_gray.lower()=='true':
        if_gray=True
    else:
        if_gray=False
    x_train, y_train= load_data(os.path.join(os.path.join(output_dir,file_name),'train'),if_gray)
    x_test, y_test=load_data(os.path.join(os.path.join(output_dir,file_name),'test'),if_gray)
    print("### loading data done.")

    x_train, x_test = x_train / 255.0, x_test / 255.0
    np.save(os.path.join(output_dir , x_train_name), x_train)
    np.save(os.path.join(output_dir , x_test_name), x_test)
    np.save(os.path.join(output_dir , y_train_name), y_train)
    np.save(os.path.join(output_dir , y_test_name), y_test)
    print("### data transform done.")

    with open(os.path.join(output_dir , 'train_test_data.txt'), 'w') as f:
        f.write(os.path.join(output_dir , x_train_name)+".npy" + ',')
        f.write(os.path.join(output_dir , x_test_name) +".npy" +',')
        f.write(os.path.join(output_dir , y_train_name) + ".npy"+',')
        f.write(os.path.join(output_dir , y_test_name)+".npy")
    print("### write train and test data name to: train_test_data.txt done.")
    report_data(100, 100, "asfasf", "test")


def parse_arguments():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(description='Kubeflow MNIST load data script')
    parser.add_argument('--data_dir', type=str, required=True, help='local file dir')
    parser.add_argument('--file_name', type=str, required=True, help='local file to be input')
    parser.add_argument('--if_gray', type=str, required=True, help='the img is gray?')
    args = parser.parse_args()
    return args


def run():
    args = parse_arguments()
    transform(args.data_dir, args.file_name,args.if_gray)
    ##transform('.', 'imgs', 'True')


if __name__ == '__main__':
    run()
